"""
Base agent implementation
"""
import json

from abc import ABC, abstractmethod
from enum import Enum, auto
from typing import Any, Dict, List, Optional

from ..core.message import Message, MessageType

from ..utils.token_count_utils import count_token

# class ContextLevel(Enum):
#     """上下文级别枚举"""
#     SESSION = auto()  # 会话级
#     TASK = auto()  # 任务级
#     GLOBAL = auto()  # 全局级


class BaseAgent(ABC):
    """Agent基类"""

    def __init__(self,
                 agent_id: str,
                 max_num_context: int = 1000,
                 use_protocol: bool = False,
                 system_prompt: str = '',
                 duty: str = ''):
        self.agent_id = agent_id
        self.model = None  # 模型对象
        self.contexts: List[Dict[str, Any]] = []
        self.max_context_length = 128000  # 默认上下文最大长度
        self.max_num_context = max_num_context  # 默认上下文最大数量
        self.use_protocol = use_protocol  # 是否使用协议
        self.sonetto_prompt = ''
        self.system_prompt = system_prompt  # 系统提示词
        self.duty = duty

    @abstractmethod
    def forward(self, message: Message) -> Message:
        """
        处理接收到的消息并返回响应
        Args:
            message: 输入消息
        Returns:
            响应消息
        """
        pass

    def add_context(self, content: str, metadata: Optional[Dict] = None):
        """
        添加上下文
        Args:
            content: 上下文内容
            level: 上下文级别
            metadata: 元数据
        """
        context_entry = {"content": content, "metadata": metadata or {}}
        self.contexts.append(context_entry)
        self._truncate_context()

    def get_context(self) -> List[Dict]:
        """获取上下文"""
        return self.contexts

    def clear_context(self):
        """
        清除上下文
        Args:
            level: 指定要清除的上下文级别，None表示清除所有级别
        """
        self.contexts.clear()

    def _truncate_context(self):
        """
        根据最大长度截断上下文
        Args:
            level: 上下文级别
        """
        current_length = len(
            str([json.loads(ctx['content']) for ctx in self.contexts]))

        while current_length > self.max_context_length and self.contexts:
            removed = self.contexts.pop(0)
            current_length -= len(str(json.loads(removed['content'])))

        while len(self.contexts) > self.max_num_context:
            removed = self.contexts.pop(0)
            current_length -= len(str(json.loads(removed['content'])))

    def create_message(self,
                       content: Any,
                       msg_type: MessageType = MessageType.TEXT,
                       receiver: Optional[str] = None,
                       metadata: Optional[Dict] = None) -> Message:
        """
        创建新消息
        Args:
            content: 消息内容
            msg_type: 消息类型
            receiver: 接收者ID
            metadata: 元数据
        Returns:
            创建的消息对象
        """
        return Message(type=msg_type,
                       content=content,
                       sender=self.agent_id,
                       receiver=receiver,
                       metadata=metadata)
